import gradio as gr
import pickle
import numpy as np
from PIL import Image

print("=" * 60)
print("正在初始化手写数字识别系统...")


# 模型加载函数
def load_knn_model(model_path='best_knn_model.pkl'):
    """加载预训练的KNN模型"""
    try:
        with open(model_path, 'rb') as model_file:
            model = pickle.load(model_file)
        print(f"✅ KNN模型加载成功 | 最佳邻居数: {model.n_neighbors}")
        return model
    except FileNotFoundError:
        print("❌ 模型文件缺失: 请确保 best_knn_model.pkl 存在")
        return None
    except Exception as error:
        print(f"❌ 模型加载异常: {error}")
        return None


# 加载机器学习模型
digit_recognizer = load_knn_model()


class ImageProcessor:
    """图像预处理类"""

    @staticmethod
    def convert_to_grayscale(image_data):
        """转换为灰度图像"""
        if isinstance(image_data, dict):
            pil_image = Image.fromarray(image_data["composite"])
        else:
            pil_image = Image.fromarray(image_data)

        return pil_image.convert('L') if pil_image.mode != 'L' else pil_image

    @staticmethod
    def resize_and_normalize(image, target_size=(8, 8)):
        """调整尺寸并归一化"""
        resized_img = image.resize(target_size, Image.Resampling.LANCZOS)
        pixel_array = np.array(resized_img)

        # 颜色反转处理
        if np.mean(pixel_array) > 128:
            pixel_array = 255 - pixel_array

        # 归一化到0-16范围
        return (pixel_array / 255.0) * 16


def process_handwriting_input(sketch_data):
    """处理手写输入数据"""
    if sketch_data is None:
        raise ValueError("请在画板上绘制数字")

    # 图像预处理流程
    gray_image = ImageProcessor.convert_to_grayscale(sketch_data)
    processed_array = ImageProcessor.resize_and_normalize(gray_image)

    return processed_array.flatten().reshape(1, -1)


def analyze_digit(sketch_input):
    """分析手写数字"""
    if digit_recognizer is None:
        return "⚠️ 系统异常: 识别模型未正确加载"

    try:
        # 数据预处理
        prepared_data = process_handwriting_input(sketch_input)

        # 数字预测
        predicted_digit = digit_recognizer.predict(prepared_data)[0]

        # 结果格式化
        result_message = f"🔢 识别数字: {predicted_digit}"

        # 置信度分析
        if hasattr(digit_recognizer, 'predict_proba'):
            probability_scores = digit_recognizer.predict_proba(prepared_data)[0]
            main_confidence = probability_scores[predicted_digit]
            result_message += f"\n🎯 可信程度: {main_confidence:.1%}"

            # 显示备选结果
            top_candidates = np.argsort(probability_scores)[-3:][::-1]
            result_message += "\n\n📈 候选数字:"
            for rank, candidate in enumerate(top_candidates):
                confidence_level = probability_scores[candidate]
                result_message += f"\n{rank + 1}. 数字 {candidate}: {confidence_level:.1%}"

        print(f"识别完成 → 数字: {predicted_digit}")
        return result_message

    except Exception as e:
        error_info = f"❌ 处理异常: {str(e)}"
        print(f"错误详情: {e}")
        return error_info


def reset_canvas():
    """重置画布"""
    return None


# 创建交互界面
app_interface = gr.Blocks(
    title="智能手写数字识别",
    theme=gr.themes.Monochrome()
)

with app_interface:
    # 页面标题
    gr.Markdown("""
    # 🖋️ 智能手写数字识别
    ### 基于KNN机器学习算法 | 支持数字0-9识别
    """)

    # 主要内容区域
    with gr.Row():
        # 输入区域
        with gr.Column():
            gr.Markdown("### ✏️ 数字绘制区")

            drawing_canvas = gr.Sketchpad(
                label="请在此绘制单个数字",
                height=280,
                width=280,
                image_mode="L",

                type="numpy"
            )

            # 操作按钮
            with gr.Row():
                analyze_button = gr.Button(
                    "🔍 识别数字",
                    variant="primary"
                )
                reset_button = gr.Button(
                    "🔄 清除画布",
                    variant="secondary"
                )

        # 输出区域
        with gr.Column():
            gr.Markdown("### 📋 识别结果")

            result_display = gr.Textbox(
                label="分析结果",
                lines=7,
                interactive=False,
                show_copy_button=True
            )

    # 使用说明
    with gr.Accordion("💡 使用说明", open=False):
        gr.Markdown("""
        **操作指南:**
        - 在左侧画板用鼠标绘制数字
        - 点击"识别数字"按钮进行分析
        - 查看右侧区域的识别结果
        - 需要重新绘制时点击"清除画布"

        **绘制建议:**
        - 尽量保持笔画连贯清晰
        - 数字大小适中，占据画板主要区域
        - 避免过于潦草或变形

        **技术信息:**
        - 识别算法: K最近邻(K-Nearest Neighbors)
        - 输入规格: 8×8像素灰度图像
        - 输出内容: 识别结果及置信度
        """)

    # 事件绑定
    analyze_button.click(
        fn=analyze_digit,
        inputs=drawing_canvas,
        outputs=result_display
    )

    reset_button.click(
        fn=reset_canvas,
        inputs=[],
        outputs=drawing_canvas
    )

# 应用程序入口
if __name__ == "__main__":
    print("=" * 60)
    print("🚀 手写数字识别系统启动中...")
    print("📊 算法类型: KNN机器学习")
    print("🎯 识别范围: 数字 0-9")
    print("=" * 60)

    if digit_recognizer is not None:
        print("✅ 系统初始化完成")
        print("🌍 启动Web服务...")
        print("📍 服务地址: http://localhost:8800")
        print("⏹️  终止服务: Ctrl + C")
        print("-" * 60)

        # 启动Web应用
        app_interface.launch(
            server_name="0.0.0.0",
            server_port=8800,
            share=False
        )
    else:
        print("❌ 系统启动失败")
        print("\n🔧 故障排查:")
        print("  1. 确认模型文件 best_knn_model.pkl 存在")
        print("  2. 检查文件权限和路径")
        print("  3. 验证模型文件完整性")